# SingMain.py - 这是一个去雾处理的示例脚本
# 实际应用中替换为真正的去雾算法

import math
import sys
import cv2 as cv
import numpy as np

def rgb2flo(img, r, g, b):
    # cols = img.cols
    # rows = img.rows
    rows, cols, c = img.shape

    cosA = math.sqrt(b * b + r * r) / (r * r + g * g + b * b)
    sinA = g / (r * r + g * g + b * b)
    cosB = b / (b * b + r * r)
    sinB = r / (b * b + r * r)

    Rr = np.array([[1., 0., 0.], [0, cosA, sinA], [0, -sinA, cosA]])
    Rg = np.array([[cosB, 0, sinB], [0, 1, 0], [-sinB, 0, cosB]])

    F = np.ones((rows, cols))
    L = np.ones((rows, cols))
    for i in range(rows):
        for j in range(cols):
            rgb = np.array([img[i, j, 2], img[i, j, 1], img[i, j, 0]])
            # xyF = rgb * Rg * Rr
            xyF1 = np.dot(rgb, Rg)
            xyF = np.dot(xyF1, Rr)
            # print(F[i][j])
            F[i][j] = xyF[2]  # A 亮度
            L[i][j] = math.sqrt(math.pow(xyF[1], 2) + math.pow(xyF[2], 2))
    return F, L


def dehaze(img, T, A):
    rows, cols, c = img.shape
    dehazeImg1 = np.ones((rows, cols, 3))

    for i in range(rows):
        for j in range(cols):
            for c in range(3):
                dehazeImg1[i][j][c] = (img[i][j][c] - A[c]) / T[i][j] + A[c]
    return dehazeImg1

def process_dehazing(input_path, output_path):
    """简单的去雾处理示例（使用暗通道先验的简化版本）"""
    # 读取图像
    # img = cv2.imread(input_path)
    
    # # 转换图像为float类型
    # img = img.astype(np.float32) / 255.0
    
    # # 计算暗通道 (min channel)
    # dark_channel = np.min(img, axis=2)
    
    # # 估计大气光
    # atmospheric_light = np.percentile(dark_channel, 99.9)
    
    # # 简化透射率估计 (实际应用中需要更复杂的计算)
    # transmission = 1 - 0.95 * (dark_channel / atmospheric_light)
    # transmission = np.clip(transmission, 0.1, 1.0)
    
    # # 恢复无雾图像
    # dehazed_image = np.zeros_like(img)
    # for i in range(3):  # 处理每个颜色通道
    #     dehazed_image[:, :, i] = (img[:, :, i] - atmospheric_light) / np.maximum(transmission, 0.1) + atmospheric_light
    
    # # 转换为8位无符号整数并保存
    # dehazed_image = np.clip(dehazed_image * 255, 0, 255).astype(np.uint8)
    # cv2.imwrite(output_path, dehazed_image)
    img = cv.imread(input_path)
    # img = cv.resize(img, (960, 480))
    img = np.array(img, np.double)  # B-G-R
    img = img / 255

    F, L = rgb2flo(img, 1, 1, 1)
    max_index = np.unravel_index(np.argmax(F, axis=None), F.shape)
    A = img[max_index]

    F, L = rgb2flo(img, A[2], A[1], A[0])
    max_index = np.unravel_index(np.argmax(L, axis=None), L.shape)
    max_L = L[max_index]
    max_index = np.unravel_index(np.argmax(F, axis=None), F.shape)
    max_F = F[max_index]

    D = 1.016 * F / max_F - 0.589 * L / max_L + 0.124313
    T = np.exp(-1.5 * D)
    A = img[max_index]
    # print(max_index)
    print(A)
    dehazeImg = dehaze(img, T, A)
    dehazeImg = cv.normalize(dehazeImg, None, 0, 255, cv.NORM_MINMAX, dtype=cv.CV_8U)
    cv.imwrite(output_path, dehazeImg)


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python SingMain.py <input_path> <output_path>")
        sys.exit(1)
        
    input_path = sys.argv[1]
    output_path = sys.argv[2]
    
    # 执行去雾处理
    process_dehazing(input_path, output_path)